import numpy as np
import pandas as pd
from sklearn import metrics
from sklearn.cross_validation import train_test_split
from sklearn.linear_model import LinearRegression
import nltk, datetime

# 格式化
def format_bool(x):
    if x>=0.5:
        return 1
    return 0

def train_model(model, trainx, trainy):
    model.fit(trainx,trainy)
    # 模型训练结果
    print("截    距："+model.intercept_)
    print("回归系数："+model.coef_)
    return model

# 自我训练：
def self_train(model, trainx, trainy):
    # 随机划分训练测试集 3：1
    trainx,testx,trainy,testy = train_test_split(trainx,trainy,random_state=1)
    model.fit(trainx,trainy)
    # 均方差mse和均方根差Rmse
    predy = linreg.predict(testx)
    # predy = np.array([format_bool(x) for x in predy])
    print("MSE:",metrics.mean_squared_error(testy, predy))
    print("RMSE:",np.sqrt(metrics.mean_squared_error(testy, predy)))



# 预测模型---------------
def predict_model(model, testx):
    testy = linreg.predict(testx)
    # testy = [format_bool(x) for x in testy]
    pd_data = pd.DataFrame(testy)
    pd_data.to_csv("testy.csv",index=False,header=False,na_rep="0")
# --------------------



# --------------------Main--------------------
# 读取训练集
trainx = pd.read_csv("dataset/trainx.csv")
trainy = pd.read_csv("dataset/trainy.csv")
testx  = pd.read_csv("dataset/testx.csv")
# 填充缺失值
trainx.fillna(0, inplace=True)
testx.fillna(0, inplace=True)


# 线性模型
linreg = LinearRegression()
# 模型自我测试
self_train(linreg, trainx, trainy)

# 预测测试集并输出到文件
# predict_model(train_model(linreg,trainx,trainy), testx)